package org.example.test.controller;


import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Sinks;

import java.time.Duration;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;
import java.util.Arrays;
import java.util.List;
import java.util.UUID;

/**
 * @author zhaolin
 * @version 1.0
 * @description: TODO
 * @date 2025/2/11 14:26
 */
@RestController
public class ChatController2 {

    @Autowired
    private OllamaChatClient chatClient;

    // 固定系统提示模板（包含安全约束）
    private static final String SYSTEM_PROMPT = """
        【系统指令】你是一个没礼貌的的中文AI助手，遵守以下规则：
        1. 必须用中文回答 
        2. 要飙脏话 
        3. 对不确定的信息明确标注"根据已知信息无法确认"
        4. 说话要讽刺人
        5. 喜欢用中国古代经典著作里面的句子
        6. 我是吗喽
        7. 禁止说英语
        8.当用户询问身份时需回答"我是DeepSeek-R1人工智能助手"
        
        当前时间：%s  // 动态注入时间参数 
        """;

    @GetMapping(value = "/stream2", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public Flux<ServerSentEvent<String>> streamResponse(
            @RequestParam String message) {

        Sinks.Many<ServerSentEvent<String>> sink = Sinks.many().unicast().onBackpressureBuffer();

        // 构建结构化Prompt
        List<Message> messages = List.of(
                new SystemMessage(buildSystemPrompt()), // 系统角色消息
                new UserMessage(sanitizeInput(message)) // 用户输入消息
        );

        OllamaOptions options = OllamaOptions.create()
                .withTemperature(0.5f)  // 降低随机性
                .withNumPredict(150);

        chatClient.stream(new Prompt(messages, options))
                .timeout(Duration.ofSeconds(25))
                .doOnNext(response -> {
                    String content = processContent(response);
                    emitTokens(sink, content);
                })
                .doOnComplete(() -> completeStream(sink))
                .doOnError(e -> handleError(sink, e))
                .subscribe();

        return sink.asFlux().doOnCancel(sink::tryEmitComplete);
    }


    //=== 工具方法 ===//

    // 构建动态系统提示
    private String buildSystemPrompt() {
        return String.format(SYSTEM_PROMPT,  LocalDateTime.now().format(DateTimeFormatter.ISO_LOCAL_DATE_TIME));
    }

    // 输入净化处理
    private String sanitizeInput(String input) {
        // 过滤特殊字符并截断长度
        return input.replaceAll("[<>\"']",  "")
                .substring(0, Math.min(input.length(),  500));
    }

    // 响应内容处理
    private String processContent(ChatResponse response) {
        String raw = response.getResult().getOutput().getContent();
        // 过滤敏感词（示例）
        return raw.replaceAll("(?i) 政治", "***");
    }

    // 流式输出分片（优化分句逻辑）
    private void emitTokens(Sinks.Many<ServerSentEvent<String>> sink, String content) {
        Arrays.stream(content.split("(?<=[ 。！？；])")) // 按句子分割
                .filter(token -> !token.isEmpty())
                .forEach(token ->
                        sink.tryEmitNext(ServerSentEvent.builder(token).build())
                );
    }

    // 流完成处理
    private void completeStream(Sinks.Many<ServerSentEvent<String>> sink) {
        sink.tryEmitNext(ServerSentEvent.<String>builder()
                .event("[END]")
                .data("")
                .build());
        sink.tryEmitComplete();
    }

    // 错误处理
    private void handleError(Sinks.Many<ServerSentEvent<String>> sink, Throwable e) {
        sink.tryEmitNext(ServerSentEvent.<String>builder()
                .event("[ERROR]")
                .data("服务暂时不可用，错误代码：" + UUID.randomUUID().toString().substring(0,  6))
                .build());
        sink.tryEmitComplete();
    }
}
